namespace Core;

public sealed class PacketAwaitable : IPacketAwaitable 
{
    private readonly PacketDispatcher _owningPacketDispatcher;
    private readonly TaskCompletionSource<BaseCommand> _taskCompletionSource = new();
    
    public PacketAwaitable(uint packetIdentifier, PacketDispatcher owningPacketDispatcher)
    {
        ArgumentNullException.ThrowIfNull(owningPacketDispatcher);
        
        Identifier = packetIdentifier;
        _owningPacketDispatcher = owningPacketDispatcher;
    }
        
    public uint Identifier { get; }
    
    public async Task<BaseCommand> WaitAsync(CancellationToken cancellationToken)
    {
        await using var register = cancellationToken.Register(() => Fail(new TimeoutException()));
        return await _taskCompletionSource.Task.ConfigureAwait(false);
    }

    public void Complete(BaseCommand command)
    {
        ArgumentNullException.ThrowIfNull(command);

        _taskCompletionSource.TrySetResult(command);
    }

    public void Fail(Exception exception)
    {
        ArgumentNullException.ThrowIfNull(exception);
            
        _taskCompletionSource.TrySetException(exception);
    }

    public void Cancel()
    {
        _taskCompletionSource.TrySetCanceled();
    }

    public void Dispose()
    {
        _owningPacketDispatcher.RemoveAwaitable(this);
    }
}